from typing import Tuple, Optional, Any, Dict, List, Callable
import numpy as np
import torch
from torch.utils.data import DataLoader
import logging

from .episodic_memory import EpisodicMemory


class MemoryIntegration:
    """
    Helper class to integrate memory components into learners
    """

    @staticmethod
    def create_memory_for_learner(learner: Any, feature_extractor: Callable) -> EpisodicMemory:
        """
        Create memory component for a learner

        Args:
            learner: Learner instance to create memory for
            feature_extractor: Feature extraction function

        Returns:
            Configured EpisodicMemory instance
        """
        # Create memory component with learner's configuration
        memory = EpisodicMemory(
            feature_extractor=feature_extractor,
            feature_dim=learner.feature_dim,
            memory_size=learner._memory_size,
            memory_per_class=learner._memory_per_class,
            fixed_memory=learner._fixed_memory,
            total_classes=learner._total_classes,
            device=learner._device,
            batch_size=learner.batch_size
        )

        # If learner already has memory, initialize the new component with it
        if hasattr(learner, '_data_memory') and len(learner._data_memory) > 0:
            memory._data_memory = learner._data_memory
            memory._targets_memory = learner._targets_memory
            memory.known_classes = learner._known_classes

        return memory

    @staticmethod
    def update_learner_from_memory(learner: Any, memory: EpisodicMemory) -> None:
        """
        Update learner's memory attributes from memory component

        Args:
            learner: Learner instance to update
            memory: EpisodicMemory instance to get data from
        """
        # Update learner's memory data
        learner._data_memory = memory._data_memory
        learner._targets_memory = memory._targets_memory

        # Update class means if available
        if hasattr(memory.class_mean_manager, 'class_means'):
            learner._class_means = memory.class_mean_manager.get_class_means()


def refactor_learner_memory(learner: Any) -> None:
    """
    Refactor a learner to use the new memory components

    Args:
        learner: Learner instance to refactor
    """
    # Define feature extractor from learner
    def feature_extractor(inputs):
        return learner.extract_token(inputs)

    # Create memory component
    memory = MemoryIntegration.create_memory_for_learner(
        learner,
        feature_extractor
    )

    # Attach memory to learner
    learner.memory = memory

    # Replace learner methods with memory-based versions
    original_build_rehearsal = learner.build_rehearsal_memory
    original_get_memory = learner._get_memory

    # Replace with new methods
    def new_build_rehearsal_memory(data_manager, per_class):
        memory.update_total_classes(learner._total_classes)
        memory.update_known_classes(learner._known_classes)
        memory.build_rehearsal_memory(data_manager, per_class)
        MemoryIntegration.update_learner_from_memory(learner, memory)

    def new_get_memory():
        return memory.get_memory()

    # Apply method replacements
    learner.build_rehearsal_memory = new_build_rehearsal_memory
    learner._get_memory = new_get_memory

    # Log the refactoring
    logging.info(f"Refactored {learner.__class__.__name__} to use component-based memory")
